import math
import torch as pt
import torch.nn as nn
from msi_utils import cumprod_exclusive
from pytorch_warp_msi import calculate_ray_sphere_intersection, \
  convert_point3d_to_normalized_theta_phi, bilinear_wrapper_5d
from torch import optim

basis_order = 3
basis_mlp = 1
basis_hidden = 64
lrelu_slope = 0.01
basis_out = 8

mlp = 4
mlp_hidden = 256 # 384
pos_level = 10
depth_level = 8

msi_layers = 8 # 16
msi_sublayers = 6 # 12

sigmoid_offset = 5


MSI_N = msi_layers * msi_sublayers

class ReluMLP(nn.Module):
  def __init__(self, mlp, hidden_dim, in_node=3, lrelu_slope=0.01, out_node=1):
    super().__init__()
    self.activation = nn.LeakyReLU(lrelu_slope)
    ls = [nn.Linear(in_node, hidden_dim)]
    ls.append(self.activation)
    for _ in range(mlp):
      ls.append(nn.Linear(hidden_dim, hidden_dim))
      ls.append(self.activation)
    ls.append(nn.Linear(hidden_dim, out_node))
    self.seq = nn.Sequential(*ls)

  def forward(self, x):
    return self.seq(x)

class NexBasis(nn.Module):
  def __init__(self, rank):
    super().__init__()
    self.seq_basis = ReluMLP(
      basis_mlp,
      basis_hidden,
      basis_order * 4,
      lrelu_slope,
      basis_out
    )
    print(f'Basis Network: {self.seq_basis}')

    # positional encoding pre compute
    self.pos_freq_viewing = pt.Tensor([0.5 * math.pi * (2 ** i) \
      for i in range(basis_order)]).view(1, 1, 1, 1, -1).cuda(rank)

  def forward(self, n_theta_phi, coeff):
    # (..., N, 2, basis_order)
    dxy = n_theta_phi[..., None] * self.pos_freq_viewing
    # (..., N, 2 * basis_order)
    dxy = dxy.view(list(dxy.shape[:-2]) + [2 * basis_order])

    dxy = pt.cat([pt.sin(dxy), pt.cos(dxy)], -1)

    basis_rgb = pt.tanh(self.seq_basis(dxy)).unsqueeze(-2) # basis rgb range: [-1, 1]
    coeff = pt.tanh(coeff).view(list(coeff.shape[:-1]) + [3,-1]) # coeff range: [-1, 1]
    return pt.sum(coeff * basis_rgb, -1)

class NexMsi(nn.Module):
  def __init__(self, dmin, dmax, msi_resolution_w, msi_resolution_h,
    theta_half_range, phi_half_range):
    super().__init__()
    self.radius = pt.linspace(dmin, dmax, MSI_N).view(-1,1,1).cuda(rank)
    msi_rgb = pt.empty([
      msi_layers,
      1,
      msi_resolution_h,
      msi_resolution_w,
      3]).uniform_(-1, 1)
    self.msi_rgb = nn.Parameter(msi_rgb)

    self.specular = NexBasis(rank)

    self.seq_k = ReluMLP(
      mlp,
      mlp_hidden,
      pos_level * 4 + depth_level * 2,
      lrelu_slope,
      1 + basis_out * 3
    )
    self.pos_freq = pt.Tensor([0.5 * math.pi * (2 ** i) \
      for i in range(pos_level)] * 2).view(1, 1, 1, 2, -1).cuda(rank)
    self.depth_freq = pt.Tensor([0.5 * math.pi * (2 ** i) \
      for i in range(depth_level)]).view(1, 1, 1, -1).cuda(rank)
    self.dmin = dmin
    self.dmax = dmax
    self.theta_half_range = theta_half_range
    self.phi_half_range = phi_half_range

    print(f'Msi Size: {self.msi_rgb.shape}')
    print(f'All combined layers: {MSI_N}')
    print(f'Min depth: {dmin}, Max depth: {dmax}')
    print(f'Layer of MLP: {mlp + 2}')
    print(f'Hidden Channel of MLP: {mlp_hidden}')
    print(f'Main Network: {self.seq_k}')

  def forward(self, pixel_rays_local, R, T, use_offset=False):
    '''
    Args:
      pixel_rays_local: [..., N, 3]; rays direction in target coordinate
      R: relative rotation, [..., 3, 3] matrices
      T: [..., 3], translations from target to source camera
        point p from target to source is accomplished via rot * p + t
      self.radius: [..., 1], hemisphere radius
    Returns:
      [..., N, C] final images
    '''

    # (..., N, 3), (..., N)
    intersections, mask = calculate_ray_sphere_intersection(
      pixel_rays_local, R, T, self.radius)

    # (..., N, 2)
    theta_phi = convert_point3d_to_normalized_theta_phi(
      intersections, self.theta_half_range, self.phi_half_range) # theta, phi in [-1, 1]

    # vxy: (..., N, 2, pos_level)
    vxy = theta_phi[..., None] * self.pos_freq
    # vxy: (..., N, 2 * pos_level)
    vxy = vxy.view(list(vxy.shape[:-2]) + [2 * pos_level])

    # vz: (..., N, depth_level)
    depth_radius = pt.linalg.norm(intersections, dim=-1, keepdim=True)
    vz = ((depth_radius - self.dmin) / (self.dmax - self.dmin) * 2 - 1) * self.depth_freq

    vxyz = pt.cat([vxy, vz], -1)
    vxyz_coords = pt.cat([pt.sin(vxyz), pt.cos(vxyz)], -1)
    # (..., N, out_node)
    out = self.seq_k(vxyz_coords)

    node = 0
    # (..., N)
    msi_a = out[..., node]
    msi_a_sig  = pt.sigmoid(msi_a - sigmoid_offset) # alpha range: [0, 1]

    node = 1
    if False:
      cof = out[..., node:]
    else:
      cof = out[::msi_sublayers, ..., node:]
      cof = pt.repeat_interleave(cof, msi_sublayers, 0)
    pixel_rays_dir = pixel_rays_local @ R.transpose(-1, -2)
    pixel_rays_dir = pt.nn.functional.normalize(pixel_rays_dir, dim=-1)
    n_theta = pt.atan2(pixel_rays_dir[...,2], pixel_rays_dir[...,0]) / math.pi
    l = pt.linalg.norm(pixel_rays_dir, dim=-1)
    n_phi = pt.acos(pixel_rays_dir[...,1] / l) / (math.pi * 0.5) - 1
    pixel_rays_theta_phi = pt.stack([n_theta, n_phi], dim=-1)
    pixel_rays_theta_phi = pixel_rays_theta_phi.expand(theta_phi.shape)
    illumination = self.specular(pixel_rays_theta_phi, coeff=cof)

    msi_rgb_sig = pt.sigmoid(self.msi_rgb) # base rgb: [0,1]
    depths = ((depth_radius - self.dmin) / (self.dmax - self.dmin) * 2 - 1)
    theta_phi_3d = pt.cat([theta_phi, depths], -1)
    # (..., N, 3)
    warped_rgb = bilinear_wrapper_5d(msi_rgb_sig, theta_phi_3d)
    # (..., N, 3)
    rgb = pt.clamp(warped_rgb + illumination, 0.0, 1.0) # clamp
    rgb[~mask] = 0
    pt.cuda.empty_cache()

    weight = cumprod_exclusive(1 - msi_a_sig)
    return pt.sum(rgb * weight[...,None] * msi_a_sig[...,None],
      dim=0, keepdim=True)

  def get_optim(self, lr):
    params1 = [p[1] for p in self.named_parameters() if 'seq_basis' in p[0]]
    params2 = [p[1] for p in self.named_parameters() if 'seq_k' in p[0]]
    return optim.Adam([
      {'params': params1, 'lr': lr * 0.01},
      {'params': params2, 'lr': lr * 0.01},
      {'params': self.msi_rgb, 'lr': lr}])
